# encoding=utf-8
"""
预测时加载PNG图片
"""
import os

import torch
import numpy as np
from pathlib import Path
from utils.Logger import get_logger

logger = get_logger()  # 配置日志


def read_npy_from_path(path_list: list, dir_path: str):
    '''
    输入10帧npy所在地址，转化为tensor给后续模型
    :param path_list:
    :return:
    '''
    numpy_arrays = []

    # 检查指定文件是否存在
    missing_files = []
    for file_name in path_list:
        file_path = os.path.join(dir_path, file_name + ".npy")
        if not os.path.exists(file_path):
            missing_files.append(file_name)

    # 如果有文件缺失，则加载最新的10个文件
    if missing_files:
        all_files = sorted(Path(dir_path).glob("*.npy"), key=os.path.getmtime, reverse=True)
        latest_files = [str(f) for f in all_files[:10]]
        logger.info("缺失文件" + str(missing_files) + ",选择加载最新的10个数据")

        # 加载最新的10个文件
        for file_path in latest_files:
            data = np.load(file_path)
            assert data.shape == (400, 400), f"Expected shape (400, 400) but got {data.shape} for {file_path}"
            numpy_arrays.append(data)
    else:
        # 加载指定的.npy文件
        for file_name in path_list:
            file_path = os.path.join(dir_path, file_name + ".npy")
            data = np.load(file_path)
            assert data.shape == (400, 400), f"Expected shape (400, 400) but got {data.shape} for {file_name}"
            numpy_arrays.append(data)
        # 将numpy数组列表转换为numpy数组，增加一个新维度以匹配(10, 1, 400, 400)
    # 然后再次增加一个批次大小的维度以匹配(1, 10, 1, 400, 400)
    numpy_stack = np.stack(numpy_arrays, axis=0)[None, :, None, :, :]
    # 如果输入不足10帧，使用最后一个帧进行补全
    if numpy_stack.shape[1] < 10:
        logger.info("数据不足，使用补全数据")
        padding_numpy = numpy_stack[:, -1:, ...]
        padding = np.tile(padding_numpy, (10 - numpy_stack.shape[1], 1, 1, 1))
        numpy_stack = np.concatenate([numpy_stack, padding], axis=1)

    # 将numpy数组转换为PyTorch张量
    tensor = torch.from_numpy(numpy_stack).float()  # 假设数据类型已经是float
    return tensor


if __name__ == '__main__':
    path = ["/Users/luofan/PyCharm/ShijiaZhuang/binData/202406050000.npy"] * 10
    tensor = read_npy_from_path(path)
    print(tensor.shape)
